import torch 
import torch.nn as nn 

if __name__ == "__main__":
    src_data = torch.load("src_data_standard.pkl")
    src_data = src_data[0:8]
    # conv1 = nn.Conv1d(in_channels=59, out_channels=16, kernel_size=5)
    src_data = src_data.permute(0, 2, 1)
    src_data = torch.unsqueeze(src_data, 1)
    print(src_data.shape)

    conv2 = nn.Conv2d(in_channels=1, out_channels=40, kernel_size=(1, 5), padding=(0,2))
    conv2_out = conv2(src_data)
    print(conv2_out.shape)

    max_pool1 = nn.MaxPool2d(kernel_size=(1, 2))
    max_pool1_out = max_pool1(conv2_out)
    print(max_pool1_out.shape)
    
    conv3 = nn.Conv2d(in_channels=40, out_channels=40, kernel_size=(1,251), padding=(0,125))
    conv3_out = conv3(max_pool1_out)
    print(conv3_out.shape)

    max_pool2 = nn.MaxPool2d(kernel_size=(1,20))
    max_pool2_out = max_pool2(conv3_out)
    print(max_pool2_out.shape)

    new = max_pool2_out.view((8, -1, 25))
    print(new.shape)
    
